import backtrader as bt

# 提供更多的 backtrader 量价分析代码示例
class MovingAverageVolumePriceStrategy(bt.Strategy):
    params = (
        ("short_ma_period", 10),
        ("long_ma_period", 50),
        ("volume_ma_period", 10),
    )

    def __init__(self):
        # 价格移动平均线
        self.short_ma = bt.indicators.SimpleMovingAverage(
            self.data.close, period=self.params.short_ma_period
        )
        self.long_ma = bt.indicators.SimpleMovingAverage(
            self.data.close, period=self.params.long_ma_period
        )
        # 成交量移动平均线
        self.volume_ma = bt.indicators.SimpleMovingAverage(
            self.data.volume, period=self.params.volume_ma_period
        )
        self.order = None

    def next(self):
        if len(self.data) > self.params.long_ma_period:
            if (
                self.short_ma[0] > self.long_ma[0]
                and self.data.volume[0] > self.volume_ma[0]
            ):
                if not self.position:
                    self.order = self.buy()
            elif (
                self.short_ma[0] < self.long_ma[0]
                and self.data.volume[0] < self.volume_ma[0]
            ):
                if self.position:
                    self.order = self.sell()

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(
                    f"BUY EXECUTED, Price: {order.executed.price}, Cost: {order.executed.value}, Size: {order.executed.size}"
                )
            elif order.issell():
                self.log(
                    f"SELL EXECUTED, Price: {order.executed.price}, Cost: {order.executed.value}, Size: {order.executed.size}"
                )
            self.order = None

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log("Order Canceled/Margin/Rejected")

    def log(self, txt):
        dt = self.datas[0].datetime.date(0)
        print(f"{dt}, {txt}")


if __name__ == "__main__":
    cerebro = bt.Cerebro()
    data = bt.feeds.GenericCSVData(
        dataname="your_data.csv",
        dtformat=("%Y-%m-%d"),
        datetime=0,
        open=1,
        high=2,
        low=3,
        close=4,
        volume=5,
        openinterest=-1,
    )
    cerebro.adddata(data)
    cerebro.addstrategy(MovingAverageVolumePriceStrategy)
    cerebro.broker.setcash(100000.0)
    cerebro.broker.setcommission(commission=0.001)
    print("Starting Portfolio Value: %.2f" % cerebro.broker.getvalue())
    cerebro.run()
    print("Final Portfolio Value: %.2f" % cerebro.broker.getvalue())
